#include "ModelConfig.h"
#include "Utils/Log.h"
#include <fstream>
#include <filesystem>


namespace AVSAnalyzer {

bool ModelConfig::loadAllConfigs() {
    try {
        namespace fs = std::filesystem;
        fs::path modelsPath("models");
        
        if (!fs::exists(modelsPath) || !fs::is_directory(modelsPath)) {
            LOGE("Models directory not found");
            return false;
        }

        bool hasValidConfig = false;
        // 遍历models目录下的所有子目录
        for (const auto& entry : fs::directory_iterator(modelsPath)) {
            if (fs::is_directory(entry)) {
                // 记录子目录名
                std::string subdirName = entry.path().filename().string();

                // 检查子目录中是否存在config.json
                fs::path configPath = entry.path() / "config.json";
                if (fs::exists(configPath)) {
                    LOGI("Loading config from: %s", configPath.string().c_str());
                    if (loadAlgorithmConfig(configPath.string(), subdirName)) {
                        hasValidConfig = true;
                    }
                }
            }
        }

        if (!hasValidConfig) {
            LOGE("No valid algorithm configs found in models directory");
            return false;
        }

        return true;
    } catch (const std::exception& e) {
        LOGE("Failed to load configs: %s", e.what());
        return false;
    }
}

bool ModelConfig::loadAlgorithmConfig(const std::string& configPath, const std::string& subdirName) {
    try {
        std::ifstream ifs(configPath);
        if (!ifs.is_open()) {
            LOGE("Failed to open algorithm config file: %s", configPath.c_str());
            return false;
        }

        Json::CharReaderBuilder builder;
        JSONCPP_STRING errs;
        Json::Value root;

        if (parseFromStream(builder, ifs, &root, &errs)) {
            AlgorithmConfig config;
            if (!parseAlgorithmConfig(root, config, subdirName)) {
                return false;
            }
            // 使用唯一的键存储配置
            std::string uniqueKey = config.type; //+ "_" + subdirName;
            mAlgorithmConfigs[uniqueKey] = config;
            LOGI("Loaded config for: %s", uniqueKey.c_str());
            return true;
        } else {
            LOGE("Failed to parse algorithm config: %s", errs.c_str());
            return false;
        }
    } catch (const std::exception& e) {
        LOGE("Failed to parse algorithm config: %s", e.what());
        return false;
    }
}

bool ModelConfig::parseAlgorithmConfig(const Json::Value& json, AlgorithmConfig& config, const std::string& subdirName) {
    try {
        // 基本信息
        config.gid = json["gid"].asInt();
        config.name = json["name"].asString();
        config.version = json["version"].asString();
        
        // Combine models root, subdir, and model file name to get the full path
        std::filesystem::path modelsRoot = "models";
        std::filesystem::path modelSubDir = subdirName;
        std::filesystem::path modelFile = json["path"].asString(); // The model file name

        // Combine paths to get the full path
        std::filesystem::path fullPath = modelsRoot / modelSubDir / modelFile;
        config.path = std::filesystem::absolute(fullPath).string();
        
        config.type = json["type"].asString();

        // 推理参数
        config.scoreThreshold = json["scoreThreshold"].asFloat();
        config.nmsThreshold = json["nmsThreshold"].asFloat();

        // 类别名称列表
        const auto& classNameArray = json["className"];
        for (const auto& name : classNameArray) {
            config.className.push_back(name.asString());
        }

        // 类别详细配置
        const auto& classesArray = json["classes"];
        for (const auto& classJson : classesArray) {
            AlgorithmClass algorithmClass;
            algorithmClass.id = classJson["id"].asInt();
            algorithmClass.name = classJson["name"].asString();
            algorithmClass.scoreThreshold = classJson["scoreThreshold"].asFloat();
            config.classes.push_back(algorithmClass);
        }

        // 算法特定参数
        const auto& paramsJson = json["algorithmParams"];
        config.algorithmParams.recognitionDistanceThresh = paramsJson["recognitionDistanceThresh"].asFloat();
        
        const auto& inputSizeArray = paramsJson["inputSize"];
        for (const auto& size : inputSizeArray) {
            config.algorithmParams.inputSize.push_back(size.asInt());
        }

        // 预处理参数
        const auto& preprocessJson = paramsJson["preprocess"];
        const auto& meanArray = preprocessJson["mean"];
        const auto& stdArray = preprocessJson["std"];
        
        for (const auto& value : meanArray) {
            config.algorithmParams.preprocess.mean.push_back(value.asFloat());
        }
        
        for (const auto& value : stdArray) {
            config.algorithmParams.preprocess.std.push_back(value.asFloat());
        }
        
        config.algorithmParams.preprocess.scale = preprocessJson["scale"].asFloat();

        return true;
    } catch (const std::exception& e) {
        LOGE("Failed to parse algorithm config: %s", e.what());
        return false;
    }
}

const AlgorithmConfig* ModelConfig::getAlgorithmConfig(const std::string& name) const {
    auto it = mAlgorithmConfigs.find(name);
    if (it != mAlgorithmConfigs.end()) {
        return &it->second;
    }
    return nullptr;
}

} // namespace AVSAnalyzer 